// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#include <pollux/dwio/common/selective_column_reader_internal.h>

namespace kumo::pollux::dwio::common {
    using dwio::common::TypeWithId;

    pollux::common::AlwaysTrue &alwaysTrue() {
        static pollux::common::AlwaysTrue alwaysTrue;
        return alwaysTrue;
    }

    dwio::common::NoHook &noHook() {
        static dwio::common::NoHook hook;
        return hook;
    }

    void ScanState::updateRawState() {
        rawState.dictionary.values =
                dictionary.values ? dictionary.values->as<void>() : nullptr;
        rawState.dictionary.numValues = dictionary.numValues;
        rawState.dictionary2.values =
                dictionary2.values ? dictionary2.values->as<void *>() : nullptr;
        rawState.dictionary2.numValues = dictionary2.numValues;
        rawState.inDictionary = inDictionary ? inDictionary->as<uint64_t>() : nullptr;
        rawState.filterCache = filterCache.data();
    }

    SelectiveColumnReader::SelectiveColumnReader(
        const TypePtr &requestedType,
        std::shared_ptr<const dwio::common::TypeWithId> fileType,
        dwio::common::FormatParams &params,
        pollux::common::ScanSpec &scanSpec)
        : memoryPool_(&params.pool()),
          requestedType_(requestedType),
          fileType_(fileType),
          formatData_(params.toFormatData(fileType, scanSpec)),
          scanSpec_(&scanSpec),
          outputRows_(memoryPool_),
          valueRows_(memoryPool_),
          outerNonNullRows_(memoryPool_),
          innerNonNullRows_(memoryPool_) {
        scanState_.rowsCopy = raw_vector<vector_size_t>(memoryPool_);
        scanState_.filterCache = raw_vector<uint8_t>(memoryPool_);
    }

    void SelectiveColumnReader::filterRowGroups(
        uint64_t rowGroupSize,
        const dwio::common::StatsContext &context,
        FormatData::FilterRowGroupsResult &result) const {
        formatData_->filterRowGroups(*scanSpec_, rowGroupSize, context, result);
    }

    const std::vector<SelectiveColumnReader *> &SelectiveColumnReader::children()
    const {
        static std::vector<SelectiveColumnReader *> empty;
        return empty;
    }

    void SelectiveColumnReader::seekTo(int64_t offset, bool readsNullsOnly) {
        if (offset == readOffset_) {
            return;
        }
        if (readOffset_ < offset) {
            if (numParentNulls_ > 0) {
                POLLUX_CHECK_LE(
                    parentNullsRecordedTo_,
                    offset,
                    "Must not seek to before parentNullsRecordedTo_");
            }
            const auto distance = offset - readOffset_ - numParentNulls_;
            numParentNulls_ = 0;
            parentNullsRecordedTo_ = 0;
            if (readsNullsOnly) {
                formatData_->skipNulls(distance, true);
            } else {
                skip(distance);
            }
            readOffset_ = offset;
        } else {
            POLLUX_FAIL(
                "Seeking backward on a ColumnReader from {} to {}",
                readOffset_,
                offset);
        }
    }

    void SelectiveColumnReader::initReturnReaderNulls(const RowSet &rows) {
        if (useBulkPath() && !scanSpec_->hasFilter()) {
            anyNulls_ = nullsInReadRange_ != nullptr;
            const bool isDense = rows.back() == rows.size() - 1;
            returnReaderNulls_ = anyNulls_ && isDense;
        } else {
            returnReaderNulls_ = false;
        }
    }

    void SelectiveColumnReader::prepareNulls(
        const RowSet &rows,
        bool hasNulls,
        int32_t extraRows) {
        if (!hasNulls) {
            anyNulls_ = false;
            return;
        }

        initReturnReaderNulls(rows);
        if (returnReaderNulls_) {
            // No need for null flags if fast path.
            return;
        }

        const auto numRows = rows.size() + extraRows;
        if (resultNulls_ && resultNulls_->unique() &&
            resultNulls_->capacity() >= bits::nbytes(numRows) + simd::kPadding) {
            resultNulls_->setSize(bits::nbytes(numRows));
        } else {
            resultNulls_ = AlignedBuffer::allocate<bool>(
                numRows + (simd::kPadding * 8), memoryPool_);
            rawResultNulls_ = resultNulls_->asMutable<uint64_t>();
        }
        anyNulls_ = false;
        // Clear whole capacity because future uses could hit uncleared data between
        // capacity() and 'numBytes'.
        simd::memset(rawResultNulls_, bits::kNotNullByte, resultNulls_->capacity());
    }

    const uint64_t *SelectiveColumnReader::shouldMoveNulls(const RowSet &rows) {
        if (rows.size() == numValues_ || !anyNulls_) {
            // Nulls will only be moved if there is a selection on values. A cast
            // alone does not move nulls.
            return nullptr;
        }
        const uint64_t *moveFrom = rawResultNulls_;
        if (returnReaderNulls_) {
            if (!(resultNulls_ && resultNulls_->unique() &&
                  resultNulls_->capacity() >= rows.size() + simd::kPadding)) {
                resultNulls_ = AlignedBuffer::allocate<bool>(
                    rows.size() + (simd::kPadding * 8), memoryPool_);
                rawResultNulls_ = resultNulls_->asMutable<uint64_t>();
            }
            moveFrom = nullsInReadRange_->as<uint64_t>();
            bits::copyBits(moveFrom, 0, rawResultNulls_, 0, rows.size());
            returnReaderNulls_ = false;
        }
        POLLUX_CHECK(resultNulls_ && resultNulls_->as<uint64_t>() == rawResultNulls_);
        POLLUX_CHECK_GT(resultNulls_->capacity() * 8, rows.size());
        return moveFrom;
    }

    void SelectiveColumnReader::setComplexNulls(
        const RowSet &rows,
        VectorPtr &result) const {
        if (!nullsInReadRange_) {
            if (result->is_nulls_writable()) {
                result->clear_nulls(0, rows.size());
            } else {
                result->reset_nulls();
            }
            return;
        }

        const bool dense = 1 + rows.back() == rows.size();
        auto &nulls = result->nulls();
        if (dense &&
            !(nulls && nulls->isMutable() &&
              nulls->capacity() >= bits::nbytes(rows.size()))) {
            result->set_nulls(nullsInReadRange_);
            return;
        }

        auto *readerNulls = nullsInReadRange_->as<uint64_t>();
        auto *resultNulls = result->mutable_nulls(rows.size())->asMutable<uint64_t>();
        if (dense) {
            bits::copyBits(readerNulls, 0, resultNulls, 0, rows.size());
            return;
        }
        for (vector_size_t i = 0; i < rows.size(); ++i) {
            bits::setBit(resultNulls, i, bits::isBitSet(readerNulls, rows[i]));
        }
    }

    void SelectiveColumnReader::getIntValues(
        const RowSet &rows,
        const TypePtr &requestedType,
        VectorPtr *result) {
        switch (requestedType->kind()) {
            case TypeKind::SMALLINT:
                switch (valueSize_) {
                    case 8:
                        getFlatValues<int64_t, int16_t>(rows, result, requestedType);
                        break;
                    case 4:
                        getFlatValues<int32_t, int16_t>(rows, result, requestedType);
                        break;
                    case 2:
                        getFlatValues<int16_t, int16_t>(rows, result, requestedType);
                        break;
                    default:
                        POLLUX_FAIL("Unsupported value size: {}", valueSize_);
                }
                break;
            case TypeKind::TINYINT:
                switch (valueSize_) {
                    case 4:
                        getFlatValues<int32_t, int8_t>(rows, result, requestedType);
                        break;
                    case 2:
                        getFlatValues<int16_t, int8_t>(rows, result, requestedType);
                        break;
                    default:
                        POLLUX_FAIL("Unsupported value size: {}", valueSize_);
                }
                break;
            case TypeKind::INTEGER:
                switch (valueSize_) {
                    case 8:
                        getFlatValues<int64_t, int32_t>(rows, result, requestedType);
                        break;
                    case 4:
                        getFlatValues<int32_t, int32_t>(rows, result, requestedType);
                        break;
                    case 2:
                        getFlatValues<int16_t, int32_t>(rows, result, requestedType);
                        break;
                    default:
                        POLLUX_FAIL("Unsupported value size: {}", valueSize_);
                }
                break;
            case TypeKind::HUGEINT:
                getFlatValues<int128_t, int128_t>(rows, result, requestedType);
                break;
            case TypeKind::BIGINT:
                switch (valueSize_) {
                    case 8:
                        getFlatValues<int64_t, int64_t>(rows, result, requestedType);
                        break;
                    case 4:
                        getFlatValues<int32_t, int64_t>(rows, result, requestedType);
                        break;
                    case 2:
                        getFlatValues<int16_t, int64_t>(rows, result, requestedType);
                        break;
                    default:
                        POLLUX_FAIL("Unsupported value size: {}", valueSize_);
                }
                break;
            default:
                POLLUX_FAIL(
                    "Not a valid type for integer reader: {}", requestedType->toString());
        }
    }

    void SelectiveColumnReader::getUnsignedIntValues(
        const RowSet &rows,
        const TypePtr &requestedType,
        VectorPtr *result) {
        switch (requestedType->kind()) {
            case TypeKind::TINYINT:
                switch (valueSize_) {
                    case 1:
                        getFlatValues<uint8_t, uint8_t>(rows, result, requestedType);
                        break;
                    case 4:
                        getFlatValues<uint32_t, uint8_t>(rows, result, requestedType);
                        break;
                    default:
                        POLLUX_FAIL("Unsupported value size: {}", valueSize_);
                }
                break;
            case TypeKind::SMALLINT:
                switch (valueSize_) {
                    case 2:
                        getFlatValues<uint16_t, uint16_t>(rows, result, requestedType);
                        break;
                    case 4:
                        getFlatValues<uint32_t, uint16_t>(rows, result, requestedType);
                        break;
                    default:
                        POLLUX_FAIL("Unsupported value size: {}", valueSize_);
                }
                break;
            case TypeKind::INTEGER:
                switch (valueSize_) {
                    case 4:
                        getFlatValues<uint32_t, uint32_t>(rows, result, requestedType);
                        break;
                    default:
                        POLLUX_FAIL("Unsupported value size: {}", valueSize_);
                }
                break;
            case TypeKind::BIGINT:
                switch (valueSize_) {
                    case 4:
                        getFlatValues<uint32_t, uint64_t>(rows, result, requestedType);
                        break;
                    case 8:
                        getFlatValues<uint64_t, uint64_t>(rows, result, requestedType);
                        break;
                    default:
                        POLLUX_FAIL("Unsupported value size: {}", valueSize_);
                }
                break;
            case TypeKind::HUGEINT:
                switch (valueSize_) {
                    case 8:
                        getFlatValues<uint64_t, uint128_t>(rows, result, requestedType);
                        break;
                    case 16:
                        getFlatValues<uint128_t, uint128_t>(rows, result, requestedType);
                        break;
                    default:
                        POLLUX_FAIL("Unsupported value size: {}", valueSize_);
                }
                break;
            default:
                POLLUX_FAIL(
                    "Not a valid type for unsigned integer reader: {}",
                    requestedType->toString());
        }
    }

    template<>
    void SelectiveColumnReader::getFlatValues<int8_t, bool>(
        const RowSet &rows,
        VectorPtr *result,
        const TypePtr &type,
        bool isFinal) {
        constexpr int32_t kWidth = ksimd::batch<int8_t>::size;
        POLLUX_CHECK_EQ(valueSize_, sizeof(int8_t));
        compactScalarValues<int8_t, int8_t>(rows, isFinal);
        auto boolValues =
                AlignedBuffer::allocate<bool>(numValues_, memoryPool_, false);
        auto rawBytes = values_->as<int8_t>();
        auto zero = ksimd::broadcast<int8_t>(0);
        if constexpr (kWidth == 32) {
            auto rawBits = boolValues->asMutable<uint32_t>();
            for (auto i = 0; i < numValues_; i += kWidth) {
                rawBits[i / kWidth] =
                        ~simd::toBitMask(zero == ksimd::load_unaligned(rawBytes + i));
            }
        } else {
            POLLUX_DCHECK_EQ(kWidth, 16);
            auto rawBits = boolValues->asMutable<uint16_t>();
            for (auto i = 0; i < numValues_; i += kWidth) {
                rawBits[i / kWidth] =
                        ~simd::toBitMask(zero == ksimd::load_unaligned(rawBytes + i));
            }
        }
        *result = std::make_shared<FlatVector<bool> >(
            memoryPool_,
            type,
            resultNulls(),
            numValues_,
            std::move(boolValues),
            std::move(stringBuffers_));
    }

    template<>
    void SelectiveColumnReader::compactScalarValues<bool, bool>(
        const RowSet &rows,
        bool isFinal) {
        if (!values_ || rows.size() == numValues_) {
            if (values_) {
                values_->setSize(bits::nbytes(numValues_));
            }
            return;
        }
        auto rawBits = reinterpret_cast<uint64_t *>(rawValues_);
        vector_size_t rowIndex = 0;
        auto nextRow = rows[rowIndex];
        auto *moveNullsFrom = shouldMoveNulls(rows);
        for (size_t i = 0; i < numValues_; i++) {
            if (outputRows_[i] < nextRow) {
                continue;
            }

            POLLUX_DCHECK_EQ(outputRows_[i], nextRow);

            bits::setBit(rawBits, rowIndex, bits::isBitSet(rawBits, i));
            if (moveNullsFrom && rowIndex != i) {
                bits::setBit(rawResultNulls_, rowIndex, bits::isBitSet(moveNullsFrom, i));
            }
            if (!isFinal) {
                outputRows_[rowIndex] = nextRow;
            }
            rowIndex++;
            if (rowIndex >= rows.size()) {
                break;
            }
            nextRow = rows[rowIndex];
        }
        numValues_ = rows.size();
        outputRows_.resize(numValues_);
        values_->setSize(bits::nbytes(numValues_));
    }

    char *SelectiveColumnReader::copyStringValue(melon::StringPiece value) {
        uint64_t size = value.size();
        if (stringBuffers_.empty() || rawStringUsed_ + size > rawStringSize_) {
            auto bytes = std::max(size, kStringBufferSize);
            BufferPtr buffer = AlignedBuffer::allocate<char>(bytes, memoryPool_);
            // Use the preferred size instead of the requested one to improve memory
            // efficiency.
            buffer->setSize(buffer->capacity());
            stringBuffers_.push_back(buffer);
            rawStringBuffer_ = buffer->asMutable<char>();
            rawStringUsed_ = 0;
            // Adjust the size downward so that the last store can take place
            // at full width.
            rawStringSize_ = buffer->capacity() - simd::kPadding;
        }
        memcpy(rawStringBuffer_ + rawStringUsed_, value.data(), size);
        auto start = rawStringUsed_;
        rawStringUsed_ += size;
        return rawStringBuffer_ + start;
    }

    void SelectiveColumnReader::addStringValue(melon::StringPiece value) {
        auto copy = copyStringValue(value);
        reinterpret_cast<StringView *>(rawValues_)[numValues_++] =
                StringView(copy, value.size());
    }

    void SelectiveColumnReader::set_nulls(BufferPtr resultNulls) {
        resultNulls_ = resultNulls;
        rawResultNulls_ = resultNulls ? resultNulls->asMutable<uint64_t>() : nullptr;
        anyNulls_ = rawResultNulls_ &&
                    !bits::isAllSet(rawResultNulls_, 0, numValues_, bits::kNotNull);
        allNull_ =
                anyNulls_ && bits::isAllSet(rawResultNulls_, 0, numValues_, bits::kNull);
        returnReaderNulls_ = false;
    }

    void SelectiveColumnReader::resetFilterCaches() {
        if (scanState_.filterCache.empty() && scanSpec_->hasFilter()) {
            scanState_.filterCache.resize(std::max<int32_t>(
                1, scanState_.dictionary.numValues + scanState_.dictionary2.numValues));
            scanState_.updateRawState();
        }
        if (!scanState_.filterCache.empty()) {
            simd::memset(
                scanState_.filterCache.data(),
                FilterResult::kUnknown,
                scanState_.filterCache.size());
        }
    }

    void SelectiveColumnReader::addParentNulls(
        int64_t firstRowInNulls,
        const uint64_t *nulls,
        const RowSet &rows) {
        const int32_t firstNullIndex =
                readOffset_ < firstRowInNulls ? 0 : readOffset_ - firstRowInNulls;
        numParentNulls_ +=
                nulls ? bits::count_nulls(nulls, firstNullIndex, rows.back() + 1) : 0;
        parentNullsRecordedTo_ = firstRowInNulls + rows.back() + 1;
    }

    void SelectiveColumnReader::addSkippedParentNulls(
        int64_t from,
        int64_t to,
        int32_t numNulls) {
        auto rowsPerRowGroup = formatData_->rowsPerRowGroup();
        if (rowsPerRowGroup.has_value() &&
            from / rowsPerRowGroup.value() >
            parentNullsRecordedTo_ / rowsPerRowGroup.value()) {
            // the new nulls are in a different row group than the last.
            parentNullsRecordedTo_ = from;
            numParentNulls_ = 0;
        }
        if (parentNullsRecordedTo_) {
            POLLUX_CHECK_EQ(parentNullsRecordedTo_, from);
        }
        numParentNulls_ += numNulls;
        parentNullsRecordedTo_ = to;
    }
} // namespace kumo::pollux::dwio::common
